package bigo;

import bigo.data.ArrayVector;
import bigo.data.MatrixMeta;
import bigo.data.RowIndex;
import bigo.data.Vector;
import bigo.lib.AgreegationDriver;
import bigo.lib.BinaryMatrixInputFormat;
import bigo.lib.BinaryMatrixOutputFormat;
import bigo.lib.InputFormat;
import bigo.lib.InputMatrix;
import bigo.lib.Mapper;
import bigo.lib.OutputFormat;
import bigo.lib.OutputMatrix;
import bigo.lib.Reducer;
import bigo.lib.TextMatrixInputFormat;
import bigo.lib.TextMatrixOutputFormat;
import java.io.IOException;
import java.util.HashMap;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileContext;
import org.apache.hadoop.fs.Options.Rename;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.util.ToolRunner;

/**
 *
 * @author songliu
 */
public class NMFAgreegationDriver extends AgreegationDriver {

    @InputMatrix
    public static MatrixMeta X, Y, Z;
    @OutputMatrix
    public static MatrixMeta O;
    @InputFormat
    public static Class inputFormat = TextMatrixInputFormat.class;
    @OutputFormat
    public static Class OutputFormat = TextMatrixOutputFormat.class;

    public NMFAgreegationDriver() {
    }

    public NMFAgreegationDriver(MatrixMeta X, MatrixMeta Y, MatrixMeta Z, MatrixMeta O) {
        NMFAgreegationDriver.X = X;
        NMFAgreegationDriver.Y = Y;
        NMFAgreegationDriver.Z = Z;
        NMFAgreegationDriver.O = O;
    }

    @Mapper
    public static class NMFMapper extends AgreegationMapper {
    }

    public static Vector agreegate(RowIndex index, HashMap<String, Vector> vectors) {
        Vector y = vectors.get(Y.path);
        Vector x = vectors.get(X.path);
        Vector z = vectors.get(Z.path);

        ArrayVector ret = new ArrayVector(x.length());
        for (Integer i : ret.keySet()) {
            ret.put(i, x.get(i) / y.get(i));
            ret.put(i, ret.get(i) * z.get(i));
        }
        return ret;
    }

    @bigo.lib.Reducer
    public static class NMFReducer extends AgreegationReducer {

        @Override
        public Vector agreegate(RowIndex index, HashMap<String, Vector> vectors) {
            return NMFAgreegationDriver.agreegate(index, vectors);
        }

    }

    public static void main(String[] args) throws Exception {
        int m = Integer.parseInt(args[3]),
                n = Integer.parseInt(args[4]);

        Configuration conf = new Configuration();
        ToolRunner.run(conf,
                new NMFAgreegationDriver(
                new MatrixMeta(args[0], m, n),
                new MatrixMeta(args[1], m, n),
                new MatrixMeta(args[2], m, n),
                new MatrixMeta("output", m, n)),
                args);
        FileContext fs = FileContext.getFileContext(conf);
        fs.delete(new Path(args[2]), true);
        fs.rename(new Path("output"), new Path(args[2]), Rename.OVERWRITE);
    }
}
